{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Model Library to Rerank the Search Results\n",
    "Boost your search results with PyMilvus's Reranker models (2.4.0+), designed to refine and prioritize your search outcomes for optimal relevance. This guide will introduce you to the integration of advanced reranker models, enhancing your search system's accuracy. Starting with a quick setup of necessary dependencies (optionally within a virtual environment), we'll cover the essentials of utilizing these models to elevate your search functionalities to new heights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install pymilvus==2.4.0\n",
    "! pip install \"pymilvus[model]\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"What event in 1956 marked the official birth of artificial intelligence as a discipline?\"\n",
    "documents = [\n",
    "    \"In 1950, Alan Turing published his seminal paper, 'Computing Machinery and Intelligence,' proposing the Turing Test as a criterion of intelligence, a foundational concept in the philosophy and development of artificial intelligence.\",\n",
    "    \"The Dartmouth Conference in 1956 is considered the birthplace of artificial intelligence as a field; here, John McCarthy and others coined the term 'artificial intelligence' and laid out its basic goals.\",\n",
    "    \"In 1951, British mathematician and computer scientist Alan Turing also developed the first program designed to play chess, demonstrating an early example of AI in game strategy.\",\n",
    "    \"The invention of the Logic Theorist by Allen Newell, Herbert A. Simon, and Cliff Shaw in 1955 marked the creation of the first true AI program, which was capable of solving logic problems, akin to proving mathematical theorems.\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BGE Rerank Function\n",
    "This rerank function uses FlagEmbedding's reranker to reorder documents according to query. For more information, please visit their [repository](https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/reranker)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3 [RerankResult(text=\"The Dartmouth Conference in 1956 is considered the birthplace of artificial intelligence as a field; here, John McCarthy and others coined the term 'artificial intelligence' and laid out its basic goals.\", score=0.9911526449170697, index=1), RerankResult(text=\"In 1950, Alan Turing published his seminal paper, 'Computing Machinery and Intelligence,' proposing the Turing Test as a criterion of intelligence, a foundational concept in the philosophy and development of artificial intelligence.\", score=0.03252822350114495, index=0), RerankResult(text='The invention of the Logic Theorist by Allen Newell, Herbert A. Simon, and Cliff Shaw in 1955 marked the creation of the first true AI program, which was capable of solving logic problems, akin to proving mathematical theorems.', score=0.006538824977235004, index=3)]\n"
     ]
    }
   ],
   "source": [
    "from pymilvus.model.reranker import BGERerankFunction\n",
    "bge_rf = BGERerankFunction(device=\"cuda:0\")\n",
    "results = bge_rf(query, documents, top_k=3)\n",
    "print(len(results), results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Voyage Rerank Function\n",
    "This rerank function uses Voyage AI's reranker to reorder documents according to query. For more information, please visit their [documentation](https://docs.voyageai.com/docs/reranker)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3 [RerankResult(text=\"The Dartmouth Conference in 1956 is considered the birthplace of artificial intelligence as a field; here, John McCarthy and others coined the term 'artificial intelligence' and laid out its basic goals.\", score=0.8984375, index=1), RerankResult(text='The invention of the Logic Theorist by Allen Newell, Herbert A. Simon, and Cliff Shaw in 1955 marked the creation of the first true AI program, which was capable of solving logic problems, akin to proving mathematical theorems.', score=0.71875, index=3), RerankResult(text=\"In 1950, Alan Turing published his seminal paper, 'Computing Machinery and Intelligence,' proposing the Turing Test as a criterion of intelligence, a foundational concept in the philosophy and development of artificial intelligence.\", score=0.6640625, index=0)]\n"
     ]
    }
   ],
   "source": [
    "from pymilvus.model.reranker import VoyageRerankFunction\n",
    "voyage_rf = VoyageRerankFunction(api_key=\"voyage-api-key\")\n",
    "results = voyage_rf(query, documents, top_k=3)\n",
    "print(len(results), results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cohere Rerank Function\n",
    "This rerank function uses Cohere's reranker to reorder documents according to query. For more information, please visit their [documentation](https://cohere.com/rerank)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3 [RerankResult(text=\"The Dartmouth Conference in 1956 is considered the birthplace of artificial intelligence as a field; here, John McCarthy and others coined the term 'artificial intelligence' and laid out its basic goals.\", score=0.99691266, index=1), RerankResult(text='The invention of the Logic Theorist by Allen Newell, Herbert A. Simon, and Cliff Shaw in 1955 marked the creation of the first true AI program, which was capable of solving logic problems, akin to proving mathematical theorems.', score=0.8578872, index=3), RerankResult(text=\"In 1950, Alan Turing published his seminal paper, 'Computing Machinery and Intelligence,' proposing the Turing Test as a criterion of intelligence, a foundational concept in the philosophy and development of artificial intelligence.\", score=0.3589146, index=0)]\n"
     ]
    }
   ],
   "source": [
    "from pymilvus.model.reranker import CohereRerankFunction\n",
    "cohere_rf = CohereRerankFunction(api_key=\"cohere-api-key\")\n",
    "results = cohere_rf(query, documents, top_k=3)\n",
    "print(len(results), results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cross-Encoder Rerank Function\n",
    "This rerank function uses Sentence-Transformers's Cross-Encoders to reorder documents according to query. For more information, please visit their [documentation](https://www.sbert.net/docs/pretrained_cross-encoders.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3 [RerankResult(text=\"The Dartmouth Conference in 1956 is considered the birthplace of artificial intelligence as a field; here, John McCarthy and others coined the term 'artificial intelligence' and laid out its basic goals.\", score=6.250533580780029, index=1), RerankResult(text=\"In 1950, Alan Turing published his seminal paper, 'Computing Machinery and Intelligence,' proposing the Turing Test as a criterion of intelligence, a foundational concept in the philosophy and development of artificial intelligence.\", score=-2.954602003097534, index=0), RerankResult(text='The invention of the Logic Theorist by Allen Newell, Herbert A. Simon, and Cliff Shaw in 1955 marked the creation of the first true AI program, which was capable of solving logic problems, akin to proving mathematical theorems.', score=-4.771515846252441, index=3)]\n"
     ]
    }
   ],
   "source": [
    "from pymilvus.model.reranker import CrossEncoderRerankFunction\n",
    "ce_rf = CrossEncoderRerankFunction(\"cross-encoder/ms-marco-MiniLM-L-6-v2\", device=\"cpu\")\n",
    "results = ce_rf(query, documents, top_k=3)\n",
    "print(len(results), results)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
